
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use("Agg")
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
import pygame
from collections import deque
import time
import multiprocessing as mp

MAX_LEN = 100

class Animation:
    def __init__(self, title="Umeko", xlabels=[""]):

        self.title = title
        self.xlabels = xlabels
        self.fig, self.ax = plt.subplots(len(xlabels),1)
        self.ln_list = []
        self.que_list = []
        self.time_list = []
        self.data_list = []
        self.start_time = time.time()
        self.count = 0
        for k, _ in enumerate(xlabels):
            if len(xlabels) == 1:
                self.ax = [self.ax]
            _ln, = self.ax[k].plot([], [])
            self.ln_list.append(_ln)
            self.que_list.append(deque([0],maxlen=MAX_LEN))
            self.time_list.append(deque([0],maxlen=MAX_LEN))
            self.data_list.append(None)
        
        plt.subplots_adjust(wspace=0, hspace=0.7)
        self.fig.set_size_inches(7, 2.5 * len(xlabels))
        
         # 将figure对象转换为pygame图像
        self.canvas = FigureCanvas(self.fig)
        self.canvas.draw()
        self.renderer = self.canvas.get_renderer()
        raw_data = self.canvas.buffer_rgba()
        self.size = self.canvas.get_width_height()
        self.surface = pygame.image.frombuffer(raw_data, self.size, "RGBA")

    # 更新数据
    def update(self, data_list):
        for k, data in enumerate(data_list) :
            self.data_list[k] = data

        for k, que in enumerate(self.que_list):
            if self.data_list[k] is not None:
                # print(que, self.time_list[k])
                que.append(self.data_list[k])
                self.time_list[k].append(self.time_list[k][-1] + 1)
                self.data_list[k] = None
                self.ax[k].set_xlim(self.time_list[k][0], self.time_list[k][-1])
                self.ax[k].set_ylim(min(que)-10, max(que)+10)
                self.ln_list[k].set_data(self.time_list[k], que)
                self.ax[k].set_title(f"{self.xlabels[k]}: {que[-1]:.2f}")    
        
        self.canvas.draw()
        raw_data = self.canvas.buffer_rgba()
        self.surface = pygame.image.frombuffer(raw_data, self.size, "RGBA")
        return self.surface

class Animation_MP:
    def __init__(self, title="Umeko", xlabels=[""], freq=1/15):
        self.title = title
        self.xlabels = xlabels
        self.parent_conn_list = []
        self.child_conn_list = []
        self.q_list = []
        self.count = 0
        self._last_update_time = time.time()
        self._freq = freq
        for _ in xlabels:
            parent_conn, child_conn = mp.Pipe()
            q = mp.Queue()
            self.parent_conn_list.append(parent_conn)
            self.child_conn_list.append(child_conn)
            self.q_list.append(q)
        # for k, i in enumerate(xlabels):
        #     process = mp.Process(target=self.run, args=(k,))
        #     process.start()
    
        
    
    def mp_update(self, data_list):
        # print(data_list)
        # for data, conn in zip(data_list, self.parent_conn_list):
        #     conn.send(data) 
        if (time.time() - self._last_update_time) > self._freq:
            self._last_update_time = time.time()
            for data, conn in zip(data_list, self.q_list):
                conn.put(data)
        
    
    def start(self):
        for k, i in enumerate(self.xlabels):
            process = mp.Process(target=self.run, args=(k,))
            process.start()
        

    def update(self, data_list):
        self.mp_update(data_list)
    
    def run(self, index):
        env = Plot_Env(self.xlabels[index])
        env.run(self.q_list[index].get)
    
    def quit(self):
        for conn in self.parent_conn_list:
            conn.send("off")



class Plot_Env():
    def __init__(self, xlabel):
        pygame.init()
        self.xlabel = xlabel
        self.fig = plt.figure()
        self.ax = self.fig.add_subplot(1, 1, 1)
        self.ln, = self.ax.plot([], [])
        self.dataque = deque([0],maxlen=MAX_LEN)
        self.timeque = deque([0],maxlen=MAX_LEN)
        self.canvas = FigureCanvas(self.fig)
        self.canvas.draw()
        self.renderer = self.canvas.get_renderer()
        raw_data = self.canvas.buffer_rgba()
        self.size = self.canvas.get_width_height()
        self.screen = pygame.display.set_mode(self.size)
        pygame.display.set_caption(f"{xlabel}")
        self.surface = pygame.image.frombuffer(raw_data, self.size, "RGBA")
        self.data = None
        self.running = True
    
    def run(self, func_data):
        self.running = True
        while self.running:
            self.update(func_data())
            self.render_curve()
            self.screen.fill((255, 255, 255))
            self.screen.blit(self.surface, (0, 0))
            pygame.display.flip()
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    self.running = False
                    pygame.quit()
    
    def render_curve(self):
        if self.data is not None:
            # print(self.data)
            if self.data == "off":
                self.running = False
                return self.surface
            self.dataque.append(self.data)
            self.timeque.append(self.timeque[-1]+1)
            self.ax.set_xlim(self.timeque[0], self.timeque[-1])
            self.ax.set_ylim(min(self.dataque)-10, max(self.dataque)+10)
            self.ln.set_data(self.timeque, self.dataque)
            self.ax.set_title(f"{self.xlabel}: {self.dataque[-1]:.2f}")
            self.canvas.draw()
            raw_data = self.canvas.buffer_rgba()
            self.surface = pygame.image.frombuffer(raw_data, self.size, "RGBA")
            self.data = None
        return self.surface
    
    def update(self, data):
        self.data = data


if __name__ == "__main__":
    
    pygame.init()
    ani = Animation_MP("Test", ["test", "test2"])
    # w, h = ani.surface.get_size()
    screen = pygame.display.set_mode((100, 100))
    pygame.display.set_caption("Matplotlib Animation in Pygame")
    k = 0
    running = True
    ani.start()
    while running:
        k += 1
        # ani.ani.event_source.refresh()
        screen.fill((255, 255, 255))
        # screen.blit(ani.surface, (0, 0))
        pygame.display.flip()
        ani.update([k, 1])
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
                pygame.quit()
        time.sleep(0.01)
